# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/25 10:10 
@Author  : ZhangShenao 
@File    : sqlalchemy.py 
@Desc    : 重写Flask-SQLAlchemy中的核心类,实现事务的自动提交
"""
from contextlib import contextmanager

from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy


class SQLAlchemy(_SQLAlchemy):
    """
    重写Flask-SQLAlchemy中的核心类,实现事务的自动提交
    """

    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e
